import datetime
import os
import traceback

from transformers import GPT2LMHeadModel, BertTokenizerFast
import torch
from parameter_config import ParameterConfig
from utils.top import top_K_P_filtering

def main():
    # 加载参数
    params = ParameterConfig()

    # 加载模型
    model = GPT2LMHeadModel.from_pretrained('./save_model/run/min_loss_model_24')
    model.to(params.device)
    model.eval()


    # 加载分词器
    tokenizer = BertTokenizerFast(params.vocab_path,
                                  unk_token="[UNK]",
                                  sep_token="[SEP]",
                                  pad_token="[PAD]",
                                  cls_token="[CLS]")
    unk_id = tokenizer.unk_token_id
    cls_id = tokenizer.cls_token_id
    sep_id = tokenizer.sep_token_id

    # 保存会话
    if params.save_history_path:
        if not os.path.exists(params.save_history_path):
            os.mkdir(params.save_history_path)
        with open(params.save_history_path + './chat_history.txt', 'a', encoding='utf-8') as f:
            f.write("[时间]{}:\n".format(datetime.datetime.now()))

    history_token = []
    # 获取用户输入
    while True:
        try:
            user_input = input("user:")

            # 将用户输入进行历史保存
            if params.save_history_path:
                with open(params.save_history_path + './chat_history.txt', 'a', encoding='utf-8') as f:
                    f.write("user:{}\n".format(user_input))

            # 输入数据文本数值化，数值张量化
            text_ids = tokenizer.encode(user_input, add_special_tokens=False)

            input_ids = [cls_id]

            input_ids.extend(text_ids)

            input_ids.append(sep_id)

            # 将输入添加到历史会话列表中
            history_token.append(text_ids)
            input_ids = [cls_id]
            # 添加历史会话
            for history_chat in history_token[-params.max_history_len:]:  # [-5:]
                input_ids.extend(history_chat)
                input_ids.append(sep_id)

            input_ids = torch.tensor(input_ids).long().to(params.device)

            # 将输入增加到二维，与模型输入一致
            input_ids = input_ids.unsqueeze(0)

            # 送入模型获取结果
            response = []
            for i in range(params.max_len):
                output = model(input_ids=input_ids)
                logits = output.logits

                # 获取最后一个结果的概率
                output_logits = logits[0, -1, :]

                # 对于已生成的结果中的每个token添加一个重复惩罚项，降低其生成概率，除等repetition_penalty，概率减小。
                for token_id in set(response):
                    output_logits[token_id] /= params.repetition_penalty

                # 将[UNK]置为无穷小，[UNK]代表模型无法识别的词汇，在生成文本时，不希望生成[UNK]，因为没有实际意义
                output_logits[unk_id] = -float('Inf')

                # 使用top_K_P_filtering函数对结果进行过滤
                logits_top_k_p = top_K_P_filtering(output_logits, top_k=params.top_k, filter_value=-float('Inf'))

                # 获取概率最大的结果
                probability = torch.softmax(logits_top_k_p, dim=-1)

                # torch.multinomial表示从候选集合中无放回地进行抽取num_samples个元素，权重越高，抽到的几率越高，返回元素的下标
                response_id = torch.multinomial(probability, num_samples=1)

                # 退出条件：碰到SEP
                if response_id == sep_id:
                    break

                # 将结果加入到模型输出列表中
                response.append(response_id)

                # 更新模型的输入，继续生成【将模型输入】加上【结果】一起送入模型，直到结束，GPT2是自回归模型，所以需要将输入加上结果一起送入模型，生成结果
                input_ids = torch.cat((input_ids, response_id.unsqueeze(0)), dim=1)

            # 将结果转换成文本
            text = tokenizer.convert_ids_to_tokens(response)
            text = "".join(text)

            # 将模型输出保存
            if params.save_history_path:
                with open(params.save_history_path + './chat_history.txt', 'a', encoding='utf-8') as f:
                    f.write("chatbot:{}\n".format(text))

            print("chatbot:", text)

        except:
            if params.save_history_path:
                f.close()
            print(traceback.format_exc())

if __name__ == '__main__':
    main()